package com.ckw.judger.utils;

import java.util.concurrent.atomic.AtomicLong;

public class TokenBucketUtils {
    private final long capacity; // 桶的容量
    private final long refillRate; // 每秒补充的令牌数
    private AtomicLong tokens; // 当前令牌数
    private long lastRefillTimestamp; // 上次补充令牌的时间戳

    public TokenBucketUtils(long capacity, long refillRate) {
        this.capacity = capacity;
        this.refillRate = refillRate;
        this.tokens = new AtomicLong(capacity);
        this.lastRefillTimestamp = System.currentTimeMillis();
    }

    /**
     * 尝试消费一个令牌
     *
     * @return 如果成功消费令牌返回 true，否则返回 false
     */
    public boolean tryConsume() {
        refill();
        long currentTokens = tokens.get();
        if (currentTokens > 0) {
            return tokens.compareAndSet(currentTokens, currentTokens - 1);
        }
        return false;
    }

    /**
     * 根据时间差补充令牌
     */
    private void refill() {
        long now = System.currentTimeMillis();
        long timeElapsed = now - lastRefillTimestamp;
        long newTokens = timeElapsed * refillRate / 1000;
        if (newTokens > 0) {
            lastRefillTimestamp = now;
            tokens.updateAndGet(current -> Math.min(capacity, current + newTokens));
        }
    }
}
